import csv
import os.path
import pandas as pd


def print_grad(model):
    """
    打印每层的梯度
    @param model: model
    @return:
    """
    for name, parms in model.named_parameters():
        try:
            print('-->name:', name, '-->grad_requirs:', parms.requires_grad, ' -->grad_value_mean:', parms.grad.mean())
        except:
            print('-->name:', name, '-->grad_requirs:', parms.requires_grad, ' -->grad_value:', parms.grad)


def log_grad(model, index, csv_path="log_grad.csv"):
    """
    记录model中每层的梯度
    @param model: model
    @param index: 本次记录的迭代次数
    @param csv_path: 路径
    @return:
    """
    grad_dict = {}
    for name, parms in model.named_parameters():
        try:
            grad_dict.update({name: float(parms.grad.mean().detach().cpu().numpy())})
        except:
            grad_dict.update({name: "None"})

    df_data = pd.DataFrame([grad_dict], index=[index])

    if not os.path.exists(csv_path):
        df_data.to_csv(csv_path)
    else:
        df_data.to_csv(csv_path, mode='a', header=False)

